from cProfile import label
from cgitb import small
from matplotlib.axes import Axes
from mpl_toolkits import mplot3d
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import matplotlib.pyplot as plt
from pyparsing import Opt
import torch as t
import torch.functional as F
import torch.nn as nn
import numpy as np
import cv2 as cv
import matplotlib.pyplot as plt
import random
from tqdm import tqdm

# #创建3d绘图区域
# ax = plt.axes(projection='3d')

# fig = plt.figure()
# #从三个维度构建
# z = np.linspace(0, 1, 100)
# x = z * np.sin(20 * z)
# y = z * np.cos(20 * z)
# #调用 ax.plot3D创建三维线图
# ax.plot3D(x, y, z, 'gray')
# ax.set_title('3D line plot')
# plt.show()

def get_circle_points(r, angle):
    x = r * np.cos(angle)
    y = r * np.sin(angle)
    return [x, y]

small_circle = np.array([get_circle_points(2, p) for p in np.linspace(0, 2 * np.pi, 50)])
big_circle = np.array([get_circle_points(6, p) for p in np.linspace(0, 2 * np.pi, 50)])

# print(get_circle_points(1, np.pi / 3))
plt.scatter(small_circle[:, 0], small_circle[:, 1], c='r')
plt.scatter(big_circle[:, 0], big_circle[:, 1], c='g')
plt.show()


class Model(nn.Module):
    
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(2, 3)
        self.fc2 = nn.Linear(3, 1)
        
    
    def forward(self, x):
        hidden_x = t.sigmoid(self.fc1(x))
        x = t.sigmoid(self.fc2(hidden_x))
        return hidden_x, x
    

model = Model()
optim = t.optim.Adam(model.parameters(), lr=0.08)
batch_size = 32

#拼接数据
all_datas = np.concatenate((small_circle, big_circle), axis=0)
all_labels = np.concatenate((np.zeros(50), np.ones(50)))
data_len = len(all_datas)
mse_loss = t.nn.BCELoss()

# plt.ion()
# fig = plt.figure()
#创建绘图区域
fig = None
loss_values = []
center_loss_values = []
for epoch in tqdm(range(100)):
    model.train()
    for i in range(int(data_len / batch_size)):
        indexes = [random.randint(0, data_len - 1) for i in range(batch_size)]
        x = all_datas[indexes]
        target = all_labels[indexes]
        x = t.from_numpy(x).float()
        target = t.from_numpy(target).float()
        hidden_x, x = model(x)
        loss = mse_loss(t.squeeze(x), target)
        optim.zero_grad()
        loss.backward()
        optim.step()
        center_loss_values.append(loss.item())
        
    if epoch % 2 == 0:
        loss_values.append(np.array(center_loss_values).mean())
        model.eval()
        x = t.from_numpy(small_circle).float()
        hidden_x, x = model(x)
        # print(hidden_x.shape)
        small_newpoints = hidden_x.detach().numpy()
        
        x = t.from_numpy(big_circle).float()
        hidden_x, x = model(x)
        big_newpoints = hidden_x.detach().numpy()

        if fig is not None:
            plt.close(fig)
        fig = plt.figure()
        ax : Axes = plt.axes(projection='3d')
        ax.scatter(small_newpoints[:, 0], small_newpoints[:, 1], small_newpoints[:, 2], c='r', label='small')
        ax.scatter(big_newpoints[:, 0], big_newpoints[:, 1], big_newpoints[:, 2], c='g', label='big')
        # plt.show()
        plt.savefig(f'./datas/{epoch}.png')
        # plt.close(fig)

fig = plt.figure()

plt.scatter([i for i in range(len(loss_values))], loss_values, c='r')
plt.show()